#!/usr/bin/env python
import sys
import time
import copy
import os.path

from Ganga.Core.GangaThread.GangaThread import GangaThread
from Ganga.Core.GangaRepository import getRegistry, RegistryKeyError
from Ganga.Core.GangaThread.MTRunner import MTRunner, Data, Algorithm

from Ganga.Utility.logging import getLogger
from AutoD3PDMakerTask import AutoD3PDMakerTask
from AutoD3PDMakerUtility import *
from SQLiteTaskStore import SQLiteTaskStore
                
class AutoD3PDMakerPostProcessor(GangaThread):
    '''
    A daemon class responsible for post-processing the completed GangaTasks
    '''
    def __init__(self, taskManager, taskHandler, name='AutoD3PDMakerPostProcessor', pollRate=10):
        GangaThread.__init__(self, name=name)
        self.taskManager = taskManager
        self.taskHandler = taskHandler
        self.pollRate    = pollRate
        self.logger      = getLogger('AutoD3PDMakerPostProcessor')

        ## MTRunner Algorithm for loading postprocess function
        class PPRunnerAlgorithm(Algorithm):

            def __init__(self, manager, handler):
                Algorithm.__init__(self)
                self.manager = manager
                self.handler = handler

            def process(self, taskObj):

                ## log the task status before post-processing
                taskObj.lastUpdate = time.time()
                taskObj.status     = AutoD3PDMakerTask.status_map['finishing']

                ## if not in batch mode, updating the log immediately as there is going to be a blocking post-process
                if not self.handler.ppBatchMode:
                    self.manager.addLogTask(copy.deepcopy(taskObj))
                
                ick, gtId_pp = self.handler.postprocessTask( taskObj )

                taskObj.lastUpdate = time.time()
                
                if ick:
                    if self.handler.ppBatchMode:
                        if gtId_pp:
                            taskObj.gtId_pp = gtId_pp
                            self.manager.addFiniTask(taskObj)
                        else:
                            taskObj.error  = 'failed to generate PP GangaTask'
                            taskObj.status = AutoD3PDMakerTask.status_map['unfinished']
                    else:
                        taskObj.status = AutoD3PDMakerTask.status_map['finished']
                else:
                    taskObj.error  = 'Post-process not fully completed'
                    taskObj.status = AutoD3PDMakerTask.status_map['unfinished']

                ## log the task status after post-processing
                self.manager.addLogTask(copy.deepcopy(taskObj))

                return ick

        self.ppRunner = MTRunner(name='AutoD3PDMakerPostProcessRunner', data=Data(collection=[]), algorithm=PPRunnerAlgorithm(self.taskManager, self.taskHandler))
        self.ppRunner.numThread = 10
        self.ppRunner.keepAlive = True

    def run(self):

        self.ppRunner.start()

        while not self.should_stop():

            now = time.time()

            self.logger.debug('Getting completed tasks from TaskManager: %s' % self.taskManager.id)

            ## retrieve completed tasks
            post_tasks = self.taskManager.getAllPostTasks()
            
            self.logger.info('Running post-process on GangaTasks: %s' % ( ','.join( map( lambda t:str(t.gtId), post_tasks ) ) ) )

            for t in post_tasks:
                self.ppRunner.addDataItem(t)

            time.sleep(self.pollRate)
            
            continue

        self.ppRunner.stop()
        self.unregister()

class AutoD3PDMakerPPMonitor(GangaThread):
    '''
    A daemon class responsible for monitoring the submitted GangaTask for post-processing
    '''
    def __init__(self, taskManager, name='AutoD3PDMakerPPMonitor', pollRate=10):
        GangaThread.__init__(self, name=name)
        self.taskManager = taskManager
        self.pollRate    = pollRate
        self.logger      = getLogger('AutoD3PDMakerPPMonitor')

    def run(self):

        while not self.should_stop():
            
            self.logger.debug('Getting finishing tasks from TaskManager: %s' % self.taskManager.id)

            ## internal list of the still-running pp tasks
            run_tasks = []
            
            ## retrieving the still-running pp-tasks
            ck_tasks = self.taskManager.getAllFiniTasks()

            self.logger.info('Checking running post-processing GangaTasks: %s' % ','.join( map( lambda t:str(t.gtId_pp), ck_tasks) ) )

            for t in ck_tasks:
                
                now = time.time()
                
                try:
                    trp = getRegistry('tasks').getProxy()

                    gt  = trp(t.gtId_pp)

                    status = gt.status
                    t.lastUpdate     = now
                    t.progress       = {}
                    t.progress['pp'] = {'all' : gt.n_all(),
                                        'run' : gt.n_status('running'),
                                        'done': gt.n_status('completed'),
                                        'fail': gt.n_status('failed')}

                    if status == AutoD3PDMakerTask.status_map['new']:
                        ## the monitor loop may get task status still 'new' when
                        ## task is still generating and submitting jobs
                        ##  - treat it as running task to prevent recreating new tasks
                        self.logger.debug('PP GangaTask %d still running' % t.gtId_pp)
                        t.status = AutoD3PDMakerTask.status_map['finishing']
                        run_tasks.append(t)

                    elif status == AutoD3PDMakerTask.status_map['pause']:
                        self.logger.debug('PP GangaTask %d paused' % t.gtId_pp)
                        t.error  = 'Associated post-processing GangaTask paused'
                        t.status = AutoD3PDMakerTask.status_map['finishing']
                        run_tasks.append(t)

                    elif status == AutoD3PDMakerTask.status_map['running']:
                        self.logger.debug('PP GangaTask %d still running' % t.gtId_pp)
                        t.error  = ''
                        t.status = AutoD3PDMakerTask.status_map['finishing']
                        run_tasks.append(t)

                    elif status == AutoD3PDMakerTask.status_map['completed']:
                        self.logger.debug('PP GangaTask %d completed' % t.gtId_pp)
                        t.status = AutoD3PDMakerTask.status_map['finished']
                    
                except RegistryKeyError:
                    ## for task no longer available in the Ganga task repository
                    self.logger.debug('GangaTask %d not found in Ganga repository' % t.gtId_pp)
                    t.status     = AutoD3PDMakerTask.status_map['unfinished']
                    t.error      = 'Associated post-processing GangaTask not found'
                    t.lastUpdate = now

                ## make a deep copy of task before putting into the log queue
                self.taskManager.addLogTask(copy.deepcopy(t))

            ## put back the still-running-tasks into the queue of running tasks
            for t in run_tasks:
                self.taskManager.addFiniTask(t)

            time.sleep(self.pollRate)

            continue

        self.unregister()

class AutoD3PDMakerReporter(GangaThread):
    '''
    A daemon class responsible for reporting current status of the D3PDMaker tasks.
    '''
    def __init__(self, taskManager, taskStoreDir='', name='AutoD3PDMakerMonitor', pollRate=10):
        GangaThread.__init__(self, name=name)
        self.taskManager    = taskManager

        if not taskStoreDir:
            taskStoreDir = os.environ['AUTOD3PDMAKER_DBDIR']
        self.taskStoreDir   = taskStoreDir

        self.pollRate       = pollRate
        self.logger         = getLogger('AutoD3PDMakerReporter')
        self.taskStoreCache = []

    def getTaskStores(self):
        ''' gets and caches the available task stores in the taskStore repasotiry '''
        for dbpath in locate( '*.db', self.taskStoreDir ):
            if dbpath not in map( lambda x:x.dbpath, self.taskStoreCache ):
                self.taskStoreCache.append( SQLiteTaskStore( dbpath ) )

        return self.taskStoreCache

    def run(self):

        while not self.should_stop():

            self.logger.debug('Generating task report using TaskManager: %s' % self.taskManager.id)

            self.taskManager.makeReport( self.getTaskStores() )

            time.sleep(self.pollRate)
        
            continue

        self.unregister()

class AutoD3PDMakerMonitor(GangaThread):
    '''
    A daemon class responsible for monitoring the submitted GangaTask
    '''
    def __init__(self, taskManager, name='AutoD3PDMakerMonitor', pollRate=10):
        GangaThread.__init__(self, name=name)
        self.taskManager = taskManager
        self.pollRate    = pollRate
        self.logger      = getLogger('AutoD3PDMakerMonitor')

    def run(self):

        while not self.should_stop():

            self.logger.debug('Getting running/paused tasks from TaskManager: %s' % self.taskManager.id)

            ## internal list of the still-running-tasks
            run_tasks = []

            ## retrieve running and paused tasks
            ck_tasks  = self.taskManager.getAllRunTasks()
            ck_tasks += self.taskManager.getAllPauseTasks()

            self.logger.info('Checking running/paused GangaTasks: %s' % ','.join( map( lambda t:str(t.gtId), ck_tasks ) ) )

            for t in ck_tasks:

                now = time.time()
                
                try:
                    trp = getRegistry('tasks').getProxy()

                    gt  = trp(t.gtId)

                    status             = gt.status
                    t.status           = status
                    t.lastUpdate       = now
                    t.progress         = {}
                    t.progress['main'] = {'all' : gt.n_all(),
                                          'run' : gt.n_status('running'),
                                          'done': gt.n_status('completed'),
                                          'fail': gt.n_status('failed')}

                    if status == AutoD3PDMakerTask.status_map['new']:
                        ## the monitor loop may get task status still 'new' when
                        ## task is still generating and submitting jobs
                        ##  - treat it as running task to prevent recreating new tasks
                        self.logger.debug('GangaTask %d still running' % t.gtId)
                        t.status = AutoD3PDMakerTask.status_map['running']
                        run_tasks.append(t)

                    elif status == AutoD3PDMakerTask.status_map['pause']:
                        self.logger.debug('GangaTask %d paused' % t.gtId)
                        t.error = 'Associated GangaTask paused'
                        self.taskManager.addPauseTask(t)

                    elif status == AutoD3PDMakerTask.status_map['running']:
                        self.logger.debug('GangaTask %d still running' % t.gtId)
                        run_tasks.append(t)

                    elif status == AutoD3PDMakerTask.status_map['completed']:
                        self.logger.debug('GangaTask %d completed' % t.gtId)
                        self.taskManager.addPostTask(t)

                except RegistryKeyError:
                    ## for task no longer available in the Ganga task repository
                    self.logger.debug('GangaTask %d not found in Ganga repository' % t.gtId)
                    t.status     = AutoD3PDMakerTask.status_map['unknown']
                    t.error      = 'Associated GangaTask not found'
                    t.lastUpdate = now

                ## make a deep copy of task before putting into the log queue
                self.taskManager.addLogTask(copy.deepcopy(t))

            ## put back the still-running-tasks into the queue of running tasks
            for t in run_tasks:
                self.taskManager.addRunTask(t)

            time.sleep(self.pollRate)

            continue

        self.unregister()

class AutoD3PDMakerSubmittor(GangaThread):
    '''
    A daemon class responsible for generating/submitting GangaTasks 
    '''
    def __init__(self, taskManager, taskHandler, name='AutoD3PDMakerSubmittor', pollRate=10):
        GangaThread.__init__(self, name=name)
        self.taskManager = taskManager
        self.taskHandler = taskHandler
        self.pollRate    = pollRate
        self.logger      = getLogger('AutoD3PDMakerSubmittor')

    def run(self):

        while not self.should_stop():

            ## check newly assigned tasks
            self.logger.debug('Getting new tasks from TaskManager: %s' % self.taskManager.id)

            t = self.taskManager.getNextNewTask()

            now = time.time()

            if t:
                ## generate GangaTask based on the task description
                gt = None

                ## catching any exception while creating a new GangaTask
                try:
                    gt = self.taskHandler.makeGangaTask(t)
                except Exception:
                    self.logger.warning('Failed to create GangaTask for task: %s' % t.id)
                    gt = None
                    exc_type, exc_value, exc_traceback = sys.exc_info()
                    t.error = msg_exception(self.logger, exc_type, exc_value, exc_traceback)

                if gt:
                    self.logger.info('Initializing new GangaTask: %d' % gt.id)

                    t.gtId = gt.id

                    ## starting GangaTasks
                    try:
                        gt.run()
                        t.status = AutoD3PDMakerTask.status_map['running']
                        t.lastUpdate = now

                        ## adding GangaTasks to q_log
                        self.taskManager.addRunTask(t)

                    except Exception:
                        self.logger.warning('Failed to start GangaTask for task: %s' % t.id)
                        exc_type, exc_value, exc_traceback = sys.exc_info()
                        t.error  = 'GangaTask startup failure: ' + msg_exception(self.logger, exc_type, exc_value, exc_traceback)
                        t.status = AutoD3PDMakerTask.status_map['failed']
                        t.lastUpdate = now

                else:
                    ## should log all kind of issues also in the task reason
                    t.error      = 'GangaTask generation failure: ' + t.error
                    t.status     = AutoD3PDMakerTask.status_map['failed']
                    t.lastUpdate = now

                ## adding GangaTasks to q_log
                self.taskManager.addLogTask(copy.deepcopy(t))
            else:
                self.logger.debug('No new task in queue')
                time.sleep( self.pollRate )

            continue

        self.unregister()

class AutoD3PDMakerDaemon(GangaThread):
    """
    main daemon program of AutoD3PDMaker
    """

    def __init__(self, taskManager, taskHandler, taskStoreDir='', name='AutoD3PDMakerDaemon', pollRate=10):
        GangaThread.__init__(self, name=name)
        self.taskManager  = taskManager
        self.taskHandler  = taskHandler

        if not taskStoreDir:
            taskStoreDir = os.environ['AUTOD3PDMAKER_DBDIR']

        ## switch to run only the reporter daemon
        self.reporter_only = False

        ## switch to turn-off the reporter daemon
        self.no_reporter   = False

        self.taskStoreDir = taskStoreDir
        self.pollRate     = pollRate
        self.logger       = getLogger('AutoD3PDMakerDaemon')

    def run(self):
        
        submittor  = None
        pprocessor = None
        monitor    = None
        monitor_pp = None
        reporter   = None

        ## daemons for the pipeline of the D3PDMaking: submittor, monitor, pprocessor, monitor_pp
        if not self.reporter_only:
            ## submittor
            self.logger.info('Starting submittor ...')
            submittor = AutoD3PDMakerSubmittor(self.taskManager, self.taskHandler)
            submittor.start()

            ## monitor
            self.logger.info('Starting monitor ...')
            monitor = AutoD3PDMakerMonitor(self.taskManager,pollRate=10)
            monitor.start()

            if self.taskHandler.ppBatchMode:
                self.logger.info('Starting PP monitor ...')
                monitor_pp = AutoD3PDMakerPPMonitor(self.taskManager,pollRate=10)
                monitor_pp.start()

            ## post-processor
            self.logger.info('Starting post-processor ...')
            pprocessor = AutoD3PDMakerPostProcessor(self.taskManager, self.taskHandler)
            pprocessor.start()

        ## daemon for generating the report
        if not self.no_reporter:
            self.logger.info('Starting task status reporter ...')
            reporter = AutoD3PDMakerReporter(self.taskManager, self.taskStoreDir, pollRate=10)
            reporter.start()
        
        try:
            ## running an infinity loop in main program
            while not self.should_stop():
                time.sleep(self.pollRate)
                self.logger.info('AutoD3PDMakerDaemon is up and alive')
        except KeyboardInterrupt:
            pass

        try:
            self.logger.info('Finishing the daemon... ')

            if reporter and reporter.isAlive():
                self.logger.info('Waiting reporter threads to finish...')
                reporter.stop()
                reporter.join()

            if pprocessor and pprocessor.isAlive():
                self.logger.info('Waiting post-processor threads to finish...')
                pprocessor.stop()
                pprocessor.join()

            if monitor_pp and monitor_pp.isAlive():
                self.logger.info('Waiting PP monitor threads to finish...')
                monitor_pp.stop()
                monitor_pp.join()

            if monitor and monitor.isAlive():
                self.logger.info('Waiting monitor threads to finish...')
                monitor.stop()
                monitor.join()

            if submittor and submittor.isAlive():
                self.logger.info('Waiting submittor threads to finish...')
                submittor.stop()
                submittor.join()

            if self.taskManager:
                self.logger.info('Dumping tasks in task manager queues...')
                self.taskManager.dumpQueuedTasks()

        finally:
            self.logger.info('AutoD3PDMakerDaemon is down')
            ## unregister the daemon from the GangaThreadPool in any case
            self.unregister()
